#!/usr/bin/python

import roslib; roslib.load_manifest('darci') 
import rospy
import sys, time, os
import math, numpy as np
import time
from geometry_msgs.msg import Wrench
from geometry_msgs.msg import WrenchStamped
from geometry_msgs.msg import Twist
from geometry_msgs.msg import TwistStamped
from m3ctrl_msgs.msg import M3JointCmd
import tf
import hrl_lib.transforms as tr
from hrl_msgs.msg import FloatArray
import time
class Lead():
   
    def __init__(self):
        self.RIGHT_ARM = 0
        self.LEFT_ARM = 1
        self.JOINT_MODE_ROS_THETA_GC = 2
        self.SMOOTHING_MODE_SLEW = 1

        self.listener = tf.TransformListener()

        #Subscribe to force readings from shared memory (should be mN and mNm units)
        force_sub_left = rospy.Subscriber('/loadx6_left_state', Wrench, self.force_cb_left)
        force_sub_right = rospy.Subscriber('/loadx6_right_state', Wrench, self.force_cb_right)
        
        #Publishers to command base and to export data for bagging
        self.base_vel_pub = rospy.Publisher('/omnibase_command', Twist) #this makes the base move
        self.base_vel_stamped_pub = rospy.Publisher('/efri/omnibase_command', TwistStamped) #for logging purposes
        self.force_l_pub = rospy.Publisher('efri/l_arm/force', WrenchStamped )
        self.force_r_pub = rospy.Publisher('efri/r_arm/force', WrenchStamped )
        #capturing rotation matrix used to transform shm force into force in torso link. 
        #For use offline if want to bag loadx6_left_state and rotate by this matrix for
        #faster force logging than this program allows (~100 Hz)
        self.force_l_rot_pub = rospy.Publisher('efri/l_arm/force_rot', FloatArray ) 
        self.force_r_rot_pub = rospy.Publisher('efri/r_arm/force_rot', FloatArray )
        
        rospy.sleep(1.0) #used to allow things to initialize and then latch properly
        
        #Variables for force sensor biasing
        self.biased_left = False
        self.biased_right = False
        self.force_left_bias = 0.0
        self.torque_left_bias = 0.0
        self.force_right_bias = np.matrix([0,0,0]).T
        self.torque_right_bias = 0.0

        self.force_z_left = 0 
        self.force_z_right = 0
        
        #Variables to hold raw, unbiased force values from shared memory
        self.force_left_raw = 0.0 
        self.force_right_raw = 0.0
        self.torque_left_raw = 0.0
        self.torque_right_raw = 0.0
        
        #Variables to store the last 10 x-velocity samples for averaging (smoothing)
        self.num_vel_samples = 10
        self.xvel_history = np.matrix(np.zeros([self.num_vel_samples,1])) 
        self.max_base_vel = 0.6 

    def force_cb_left(self,data):
        self.force_z_left = data.force.z
        if np.abs(self.force_z_left) < 4000: 
            self.force_z_left = 0.0

        #Divide by 1000 since raw values should be mN and mNm
        self.force_left_raw = np.matrix([data.force.x, data.force.y, data.force.z]).T / 1000.0
        self.torque_left_raw = np.matrix([data.torque.x, data.torque.y, data.torque.z]).T / 1000.0

    def force_cb_right(self,data):
        self.force_z_right = data.force.z
        if np.abs(self.force_z_right) < 4000: 
            self.force_z_right = 0.0
            
        self.force_right_raw = np.matrix([data.force.x, data.force.y, data.force.z]).T / 1000.0
        self.torque_right_raw = np.matrix([data.torque.x, data.torque.y, data.torque.z]).T / 1000.0
    
    def bias_force_sensor(self,arm):
        #arm is 'l' or 'r'
        if arm == 'l':
            self.force_left_bias = self.force_left_raw
            self.torque_left_bias = self.torque_left_raw 
            self.biased_left = True
            print 'Left force bias is: ', self.force_left_bias
        elif arm == 'r':
            self.force_right_bias = self.force_right_raw
            self.torque_right_bias = self.torque_right_raw 
            self.biased_right = True
            print 'Right force bias is: ', self.force_right_bias

    def get_force(self,arm,bias=True,torso_frame=True,type='Wrench'):
        #arm is 'l' or 'r'

        if bias:
            if arm == 'l' and self.biased_left == True:
                force = self.force_left_raw - self.force_left_bias
                torque = self.torque_left_raw - self.torque_left_bias
            elif arm == 'r' and self.biased_right == True:
                force = self.force_right_raw - self.force_right_bias
                torque = self.torque_right_raw - self.torque_right_bias
            else:
                print arm, ' arm is not biased. Bias using bias_force_sensor(arm). Using raw values instead.'
                if arm == 'l':
                    force = self.force_left_raw
                    torque = self.torque_left_raw
                elif arm == 'r':
                    force = self.force_right_raw
                    torque = self.torque_right_raw
        else:
            if arm == 'l':
                force = self.force_left_raw
                torque = self.torque_left_raw
            elif arm == 'r':
                force = self.force_right_raw
                torque = self.torque_right_raw

        if torso_frame:
            if arm == 'l':
                (trans,rot) = self.listener.lookupTransform('/torso_lift_link','/end_effector_LEFT', rospy.Time(0))
            elif arm == 'r':
                (trans,rot) = self.listener.lookupTransform('/torso_lift_link','/end_effector_RIGHT', rospy.Time(0))
     
            #rotate force into tf frame for end effector
            R = np.matrix('1, 0, 0; 0, 0, -1; 0, 1, 0')
            P = np.matrix('-1, 0, 0;0,1,0 ;0,0,-1 ')
            force = R*P*force
            rot_mat = tr.quaternion_to_matrix(rot)  #convert to 3x3 numpy rotation matrix
            trans_vec = np.matrix(trans).T #put into numpy column vector
            force = rot_mat*force #used this since need to transform vector not just a point. Else will result in trans_vec added to magnitude of force
            torque = rot_mat*torque + trans_vec
        
        if type == 'Wrench':
            data = Wrench()
            data.force.x = force[0]
            data.force.y = force[1]
            data.force.z = force[2]
            data.torque.x = torque[0]
            data.torque.y = torque[1]
            data.torque.z = torque[2]
            return data,rot_mat
        elif type == 'np': #return numpy
            return (force, torque),rot_mat
            
    def run(self):
        #Instantiate variables that will be published
        base_vel = Twist() #for base velocity command
        base_vel_stamped = TwistStamped() #for base velocity command
        force_l = WrenchStamped()
        force_r = WrenchStamped()
        force_l_rot = FloatArray()
        force_r_rot = FloatArray()


        force_l.wrench, force_l_rot.data = self.get_force('l')  
        force_r.wrench, force_r_rot.data = self.get_force('r')
        force_sum = force_l.wrench.force.x + force_r.wrench.force.x
    
        if np.abs(force_sum) < 1.0:
            force_sum = [0.0]

        base_vel_sign = np.sign(force_sum)
        base_vel.linear.x = base_vel_sign*np.min([.04*np.float(np.abs(force_sum)), self.max_base_vel]) 
        base_vel.linear.y = 0.0
        base_vel.linear.z = 0.0
        base_vel.angular.x = 0.0
        base_vel.angular.y = 0.0
        base_vel.angular.z = 0.0
        self.xvel_history[0:(self.num_vel_samples-1)] = self.xvel_history[1:self.num_vel_samples]
        self.xvel_history[(self.num_vel_samples-1)] = base_vel.linear.x 
        base_vel.linear.x = np.mean(self.xvel_history.tolist())   
        self.xvel_history[(self.num_vel_samples-1)] = base_vel.linear.x

        base_vel_stamped.twist = base_vel
        
        #Publish base command and other data for bagging
        curr_time = rospy.Time.now()
        force_l.header.stamp  = curr_time
        force_r.header.stamp  = curr_time
        force_l_rot.header.stamp  = curr_time
        force_r_rot.header.stamp  = curr_time
        base_vel_stamped.header.stamp = curr_time
        
        self.base_vel_pub.publish(base_vel)
        self.base_vel_stamped_pub.publish(base_vel_stamped)
        self.force_l_pub.publish(force_l)
        self.force_r_pub.publish(force_r)

if __name__ == '__main__':
    
    rospy.init_node( 'lead_node', anonymous = True )
    leader = Lead()
    rospy.sleep(8.0) #allow time for things to catch. else will get zero force readings
    print 'done sleeping. biasing force sensors...'
    leader.bias_force_sensor('r')
    leader.bias_force_sensor('l')
    s = raw_input('Done biasing force sensors. Do you want to begin force-Vicon synchronization using sinusoidal input ? [y/n] ')
    if s == 'y':
        print 'Collecting force data for force-Vicon calibration...'
        start_time = time.time()
        collect_time = 0.0
        while collect_time < 20.0:
            curr_ros_time = rospy.Time.now()
            curr_time = time.time()

            collect_time = curr_time - start_time
        
            force_l = WrenchStamped()
            force_r = WrenchStamped()
            force_l_rot = FloatArray()
            force_r_rot = FloatArray()
        
            force_l.wrench, force_l_rot.data = leader.get_force('l')  
            force_r.wrench, force_r_rot.data = leader.get_force('r')
            force_l.header.stamp  = curr_ros_time
            force_r.header.stamp  = curr_ros_time
            force_l_rot.header.stamp  = curr_ros_time
            force_r_rot.header.stamp  = curr_ros_time
        
            leader.force_l_pub.publish(force_l)
            leader.force_r_pub.publish(force_r)
        
            rospy.sleep(0.005)
        
        print 'Done collecting force data.'
    print 'Beginning run() program.'
    while not rospy.is_shutdown():
        leader.run()
        rospy.sleep(0.005)
        continue
    rospy.spin()

